// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <pollux/common/process/thread_local_registry.h>

#include <melon/synchronization/baton.h>
#include <melon/synchronization/latch.h>
#include <gtest/gtest.h>

#include <atomic>

namespace kumo::pollux::process {
namespace {

template <typename Tag>
class TestObject {
 public:
  static std::atomic_int& count() {
    static std::atomic_int value;
    return value;
  }

  TestObject() : threadId_(std::this_thread::get_id()) {
    ++count();
  }

  ~TestObject() {
    --count();
  }

  std::thread::id threadId() const {
    return threadId_;
  }

 private:
  const std::thread::id threadId_;
};

TEST(ThreadLocalRegistryTest, basic) {
  struct Tag {};
  using T = TestObject<Tag>;
  ASSERT_EQ(T::count(), 0);
  auto registry = std::make_shared<ThreadLocalRegistry<T>>();
  registry->forAllValues([](const T&) { FAIL(); });
  thread_local ThreadLocalRegistry<T>::Reference ref(registry);
  const T* object = ref.withValue([](const T& x) {
    EXPECT_EQ(T::count(), 1);
    return &x;
  });
  ASSERT_EQ(object->threadId(), std::this_thread::get_id());
  ref.withValue([&](const T& x) { ASSERT_EQ(&x, object); });
  int count = 0;
  registry->forAllValues([&](const T& x) {
    ++count;
    ASSERT_EQ(x.threadId(), std::this_thread::get_id());
  });
  ASSERT_EQ(count, 1);
  ASSERT_EQ(T::count(), 1);
}

TEST(ThreadLocalRegistryTest, multiThread) {
  struct Tag {};
  using T = TestObject<Tag>;
  ASSERT_EQ(T::count(), 0);
  auto registry = std::make_shared<ThreadLocalRegistry<T>>();
  constexpr int kNumThreads = 7;
  std::vector<std::thread> threads;
  melon::Latch latch(kNumThreads);
  melon::Baton<> batons[kNumThreads];
  const T* objects[kNumThreads];
  for (int i = 0; i < kNumThreads; ++i) {
    threads.emplace_back([&, i] {
      thread_local ThreadLocalRegistry<T>::Reference ref(registry);
      objects[i] = ref.withValue([](const T& x) { return &x; });
      latch.count_down();
      batons[i].wait();
    });
  }
  latch.wait();
  std::vector<int> indices;
  registry->forAllValues([&](const T& x) {
    auto it = std::find(std::begin(objects), std::end(objects), &x);
    indices.push_back(it - std::begin(objects));
  });
  ASSERT_EQ(indices.size(), kNumThreads);
  std::sort(indices.begin(), indices.end());
  for (int i = 0; i < kNumThreads; ++i) {
    ASSERT_EQ(indices[i], i);
    ASSERT_EQ(objects[i]->threadId(), threads[i].get_id());
    ASSERT_EQ(T::count(), kNumThreads - i);
    batons[i].post();
    threads[i].join();
  }
  ASSERT_EQ(T::count(), 0);
  registry->forAllValues([](const T&) { FAIL(); });
}

} // namespace
} // namespace kumo::pollux::process
